#include <opencv2/opencv.hpp>
#include <iostream>
#include <string>
#include <chrono>
#include <ctime>
#include <filesystem>
#include <sstream>
#include "camera_class.h"

namespace fs = std::filesystem;

// 创建带时间戳的输出目录
std::string createOutputDirectory(const std::string &baseDir = "dataset")
{
    auto now = std::chrono::system_clock::now();
    std::time_t now_time = std::chrono::system_clock::to_time_t(now);
    struct tm timeinfo;

    // 跨平台兼容的本地时间获取
#ifdef _WIN32
    localtime_s(&timeinfo, &now_time);
#else
    localtime_r(&now_time, &timeinfo);
#endif

    char timestamp[20];
    std::strftime(timestamp, sizeof(timestamp), "%Y%m%d_%H%M%S", &timeinfo);

    std::string outputDir = baseDir + "/capture_" + timestamp;
    fs::create_directories(outputDir);

    return outputDir;
}

// 获取当前时间戳字符串
std::string getCurrentTimestamp(bool withMs = false)
{
    auto now = std::chrono::system_clock::now();
    std::time_t now_time = std::chrono::system_clock::to_time_t(now);
    struct tm timeinfo;

#ifdef _WIN32
    localtime_s(&timeinfo, &now_time);
#else
    localtime_r(&now_time, &timeinfo);
#endif

    char timestamp[30];
    if (withMs)
    {
        std::strftime(timestamp, sizeof(timestamp), "%Y%m%d_%H%M%S", &timeinfo);

        // 添加毫秒
        auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                      std::chrono::system_clock::now().time_since_epoch()) %
                  1000;

        std::stringstream ss;
        ss << timestamp << "_" << std::setw(3) << std::setfill('0') << ms.count();
        return ss.str();
    }
    else
    {
        std::strftime(timestamp, sizeof(timestamp), "%Y%m%d_%H%M%S", &timeinfo);
        return std::string(timestamp);
    }
}

void captureDataset()
{
    camera cam;
    cv::Mat frame;
    cam.start_cam();

    cam.optimizeForDetection();

    // 创建输出目录
    std::string outputDir = createOutputDirectory();
    std::cout << "Saving images to: " << outputDir << std::endl;

    // 拍摄参数
    int frameCount = 0;
    const int maxImages = 700;     // 最大拍摄数量
    const int captureInterval = 1; // 拍摄间隔（秒）
    auto lastCaptureTime = std::chrono::steady_clock::now();

    while (frameCount < maxImages)
    {
        // 读取帧
        cam.get_pic(&frame);
        if (frame.empty())
        {
            std::cerr << "Error: Failed to capture frame." << std::endl;
            break;
        }
        cv::Mat displayframe = frame.clone();

        // 显示拍摄计数和提示
        std::string statusText = "Auto-capturing: " + std::to_string(frameCount) +
                                 " / " + std::to_string(maxImages) + " (1 image/sec)";
        std::string nextCaptureText;

        // 计算下次拍摄剩余时间
        auto currentTime = std::chrono::steady_clock::now();
        std::chrono::duration<double> elapsed = currentTime - lastCaptureTime;
        double remaining = captureInterval - elapsed.count();
        if (remaining < 0)
            remaining = 0;

        nextCaptureText = "Next capture in: " + std::to_string(remaining).substr(0, 4) + "s";
        std::string instructionText = "Press 'q' to quit, ' ' to manually capture";

        // 绘制文本到帧上
        cv::putText(displayframe, statusText, cv::Point(10, 30),
                    cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 0), 2);
        cv::putText(displayframe, nextCaptureText, cv::Point(10, 60),
                    cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 0), 2);
        cv::putText(displayframe, instructionText, cv::Point(10, 90),
                    cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 0), 2);

        // 显示预览
        cv::imshow("Camera Preview", displayframe);

        // 检查是否到了拍摄时间
        if (elapsed.count() >= captureInterval)
        {
            std::string timestamp = getCurrentTimestamp();
            std::string filename = outputDir + "/image_" + timestamp + ".jpg";

            if (cv::imwrite(filename, frame))
            {
                frameCount++;
                lastCaptureTime = currentTime;
                std::cout << "Auto-captured image " << frameCount << ": " << filename << std::endl;
            }
            else
            {
                std::cerr << "Error saving image: " << filename << std::endl;
            }
        }

        // 键盘控制
        char key = cv::waitKey(1) & 0xFF;
        if (key == 'q') // 按q退出
        {
            std::cout << "Capture stopped by user." << std::endl;
            break;
        }
        else if (key == ' ') // 按空格手动拍摄
        {
            std::string timestamp = getCurrentTimestamp(true);
            std::string filename = outputDir + "/manual_" + std::to_string(frameCount) + "_" + timestamp + ".jpg";

            if (cv::imwrite(filename, frame))
            {
                frameCount++;
                std::cout << "Manually captured image " << frameCount << ": " << filename << std::endl;
            }
            else
            {
                std::cerr << "Error saving image: " << filename << std::endl;
            }
        }
    }

    std::cout << "Capture completed. Total images: " << frameCount << std::endl;

    // 清理资源
    frame.release();
    cv::destroyAllWindows();
}

int main()
{
    std::cout << "Starting dataset capture..." << std::endl;
    std::cout << "Auto-capturing 1 image per second" << std::endl;
    std::cout << "Press 'q' to quit, ' ' to manually capture" << std::endl;

    captureDataset();

    return 0;
}